# Copyright (c) 2023 Presto Labs Pte. Ltd.
# Author: leon

import datetime
import logging
import signal
from typing import List
import sys
from functools import partial

import tornado.httpserver
import tornado.websocket
import tornado.ioloop
import tornado.web
from absl import app, flags

from coin.base.datetime_util import parse_end_time
from coin.exchange.uniswap_v3 import config
from coin.exchange.uniswap_v3.ws.subscriber import UniswapV3WsSubscriber
from coin.flow.subscriber import from_queue
from coin.flow.simple_queue import SimpleQueue

FLAGS = flags.FLAGS


class WsHandler(tornado.websocket.WebSocketHandler):
    clients = []

    def open(self):
        logging.info('New connection.')
        self.clients.append(self)

    def on_message(self, message):
        logging.info('Message received: %s' % message)

    def on_close(self):
        logging.info('Connection closed.')
        self.clients.remove(self)

    @classmethod
    def write_to_clients(cls, data):
        for client in cls.clients:
            client.write_message(data)


class RawFeedHandler(object):
    _clients: List[WsHandler]

    def __init__(self, clients: List[WsHandler]):
        self._clients = clients

    def on_book(self, record, queue_data, topic_data):
        for client in self._clients:
            try:
                client.write_message(record.data)
            except tornado.websocket.WebSocketClosedError:
                self._clients.remove(client)

    def on_trade(self, record, queue_data, topic_data):
        for client in self._clients:
            try:
                client.write_message(record.data)
            except tornado.websocket.WebSocketClosedError:
                self._clients.remove(client)


def main(argv):
    queue = SimpleQueue('uniswap-v3-forwarder')
    ioloop = tornado.ioloop.IOLoop.current()
    logger = logging.getLogger('uniswapv3_forwarder')

    sub_dict = {}

    def on_disconnect(worker_id: int):
        logging.info('Disconnected. Restarting in 3 sec.')
        ioloop.add_timeout(datetime.timedelta(seconds=3), partial(init_sub, worker_id))

    def init_sub(worker_id: int):
        sub = sub_dict.get(worker_id, None)
        if sub is not None:
            sub.stop()
        group_id = 'a'
        port = 9544 + worker_id
        sub = UniswapV3WsSubscriber(
            queue,
            group_id,
            str(worker_id),
            config.group_symbol[group_id],
            provider=f'ws://127.0.0.1:{port}/ws',
            blockchain_name='ethereum',
            close_callback=partial(on_disconnect, worker_id),
            logger=logger)
        sub.start(ioloop)
        sub_dict[worker_id] = sub

    init_sub(1)
    init_sub(2)
    flow_sub = from_queue(queue)
    raw_feed_handler = RawFeedHandler(WsHandler.clients)
    flow_sub.subscribe('uniswapv3_ws_depth_*%d' % 1, raw_feed_handler.on_book)
    flow_sub.subscribe('uniswapv3_ws_trade_*%d' % 1, raw_feed_handler.on_trade)
    flow_sub.subscribe('uniswapv3_ws_depth_*%d' % 2, raw_feed_handler.on_book)
    flow_sub.subscribe('uniswapv3_ws_trade_*%d' % 2, raw_feed_handler.on_trade)

    def sig_handler(self, sig=None, frame=None):
        ioloop.add_callback_from_signal(ioloop.stop)

    end_time = parse_end_time('23:59:50')
    ioloop.add_timeout(end_time - datetime.datetime.now(), ioloop.stop)
    signal.signal(signal.SIGTERM, sig_handler)
    signal.signal(signal.SIGINT, sig_handler)

    application = tornado.web.Application([
        (r'/ws', WsHandler),
    ])
    http_server = tornado.httpserver.HTTPServer(application)
    http_server.listen(9544)

    ioloop.start()


if __name__ == '__main__':
    logging.basicConfig(level='INFO', format='%(levelname)s %(asctime)s %(name)s] %(message)s')
    sys.exit(main(sys.argv))
